import cv2
import mediapipe as mp
import time
import numpy as np
import threading
from queue import Queue, Empty

# 全局配置（集中管理参数）
CONFIG = {
    "cam_width": 800,  # 摄像头分辨率
    "cam_height": 600,
    "win_width": 800,  # 显示窗口大小
    "win_height": 600,
    "queue_size": 5,  # 队列大小
    "fps_window": 10,  # FPS计算窗口
    "exit_timeout": 2  # 退出超时时间（秒）
}

# 全局变量（线程通信+状态）
frame_queue = Queue(maxsize=CONFIG["queue_size"])
result_queue = Queue(maxsize=CONFIG["queue_size"])
exit_flag = False  # 全局退出标志
face_mesh = None


def init_mediapipe():
    """初始化MediaPipe人脸网格模型"""
    global face_mesh
    mp_face_mesh = mp.solutions.face_mesh
    face_mesh = mp_face_mesh.FaceMesh(
        max_num_faces=1,
        refine_landmarks=True,
        min_detection_confidence=0.6,
        min_tracking_confidence=0.6,
        static_image_mode=False
    )


def capture_frames():
    """线程1：捕获摄像头帧"""
    global exit_flag
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, CONFIG["cam_width"])
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CONFIG["cam_height"])

    if not cap.isOpened():
        print("❌ 无法打开摄像头")
        exit_flag = True
        return

    time.sleep(1)  # 等待摄像头初始化
    while not exit_flag:
        ret, frame = cap.read()
        if not ret:
            time.sleep(0.1)
            continue

        # 队列满时直接丢弃新帧（避免阻塞捕获）
        if not frame_queue.full():
            frame_queue.put(frame)

    # 退出前清空队列，避免处理线程阻塞
    while not frame_queue.empty():
        try:
            frame_queue.get_nowait()
            frame_queue.task_done()
        except Empty:
            break

    cap.release()
    print("✅ 捕获线程退出")


def process_frames():
    """线程2：人脸检测+特征点绘制"""
    global exit_flag
    mp_drawing = mp.solutions.drawing_utils
    mp_styles = mp.solutions.drawing_styles
    mp_face_mesh = mp.solutions.face_mesh

    while not exit_flag:
        try:
            # 超时1秒获取帧，避免无限阻塞
            frame = frame_queue.get(timeout=0.5)
            frame_queue.task_done()

            # 预处理+检测
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = face_mesh.process(np.ascontiguousarray(frame_rgb))

            # 绘制特征点
            if results.multi_face_landmarks:
                for landmarks in results.multi_face_landmarks:
                    mp_drawing.draw_landmarks(
                        frame, landmarks, mp_face_mesh.FACEMESH_TESSELATION,
                        None, mp_styles.get_default_face_mesh_tesselation_style()
                    )
                    mp_drawing.draw_landmarks(
                        frame, landmarks, mp_face_mesh.FACEMESH_CONTOURS,
                        None, mp_styles.get_default_face_mesh_contours_style()
                    )

            # 翻转画面+计算处理时间
            frame = cv2.flip(frame, 1)
            process_duration = max(0.001, time.time() - time.time() + 0.001)  # 避免0值

            # 结果放入队列（满时丢弃）
            if not result_queue.full():
                result_queue.put((frame, process_duration))

        except Empty:
            continue
        except Exception as e:
            if not exit_flag:
                print(f"⚠️  处理错误: {str(e)}")
                time.sleep(0.1)

    # 退出前清空队列
    while not result_queue.empty():
        try:
            result_queue.get_nowait()
            result_queue.task_done()
        except Empty:
            break

    print("✅ 处理线程退出")


def show_frames():
    """主线程：显示画面+计算FPS"""
    global exit_flag
    # 初始化窗口
    cv2.namedWindow('Face Detection', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Face Detection', CONFIG["win_width"], CONFIG["win_height"])

    fps = 0.0
    frame_times = []

    while not exit_flag:
        try:
            # 获取处理后帧（超时0.5秒）
            frame, duration = result_queue.get(timeout=0.5)
            result_queue.task_done()

            # 计算滑动平均FPS
            frame_times.append(duration)
            if len(frame_times) > CONFIG["fps_window"]:
                frame_times.pop(0)
            fps = 1.0 / (sum(frame_times) / len(frame_times)) if frame_times else 0.0

            # 绘制FPS
            cv2.putText(
                frame, f"FPS: {fps:.1f}", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2, cv2.LINE_AA
            )

            # 显示画面
            cv2.imshow('Face Detection', frame)
            # 检测退出按键（ESC/Q）
            if cv2.waitKey(1) & 0xFF in [27, ord('q')]:
                exit_flag = True
                break

        except Empty:
            # 队列空时显示黑色画面（避免闪烁）
            cv2.imshow('Face Detection', np.zeros((CONFIG["cam_height"], CONFIG["cam_width"], 3), np.uint8))
            continue


def safe_queue_join(queue, timeout):
    """安全的队列等待：超时后强制退出，避免无限阻塞"""
    start_time = time.time()
    while time.time() - start_time < timeout:
        if queue.unfinished_tasks == 0:
            return True  # 所有任务完成
        time.sleep(0.1)  # 短暂等待
    # 超时后清空队列
    while not queue.empty():
        try:
            queue.get_nowait()
            queue.task_done()
        except Empty:
            break
    return False


def main():
    global exit_flag
    try:
        # 初始化模型+启动线程
        init_mediapipe()
        t_capture = threading.Thread(target=capture_frames, daemon=True)
        t_process = threading.Thread(target=process_frames, daemon=True)
        t_capture.start()
        t_process.start()

        # 启动显示逻辑
        show_frames()

    finally:
        # 1. 触发所有线程退出
        exit_flag = True

        # 2. 等待线程结束（超时处理）
        t_capture.join(timeout=CONFIG["exit_timeout"])
        t_process.join(timeout=CONFIG["exit_timeout"])

        # 3. 安全清空队列（避免join阻塞）
        safe_queue_join(frame_queue, CONFIG["exit_timeout"])
        safe_queue_join(result_queue, CONFIG["exit_timeout"])

        # 4. 释放所有资源
        if face_mesh:
            face_mesh.close()
        cv2.destroyAllWindows()
        print("✅ 程序完全退出")


if __name__ == "__main__":
    main()